{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/Goekdeniz-Guelmez/JOSIE-R1-Zero.git\n",
    "%cd JOSIE-R1-Zero\n",
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "from trl import GRPOConfig, GRPOTrainer\n",
    "\n",
    "import torch\n",
    "\n",
    "from dataset_hanlder import get_gsm8k_question\n",
    "from reward_functions import (\n",
    "    xmlcount_reward_func,\n",
    "    soft_format_reward_func,\n",
    "    strict_format_reward_func,\n",
    "    int_reward_func,\n",
    "    accuracy_reward_func\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "model_name = \"Qwen/Qwen2.5-1.5B\"\n",
    "output_dir = \"grpo_outputs/qwen_1.5b_zero\"\n",
    "run_name = \"Qwen-1.5B-Zero\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "training_args = GRPOConfig(\n",
    "    output_dir=output_dir,\n",
    "    run_name=run_name,\n",
    "    learning_rate=5e-6,\n",
    "    adam_beta1 = 0.9,\n",
    "    adam_beta2 = 0.99,\n",
    "    weight_decay = 0.1,\n",
    "    warmup_ratio = 0.1,\n",
    "    lr_scheduler_type='cosine',\n",
    "    logging_steps=1,\n",
    "    bf16=False,\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    num_generations=4,\n",
    "    max_prompt_length=256,\n",
    "    max_completion_length=786,\n",
    "    num_train_epochs=1,\n",
    "    save_steps=100,\n",
    "    max_grad_norm=0.1,\n",
    "    gradient_checkpointing=True,\n",
    "    report_to=\"wandb\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "dataset = get_gsm8k_question()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "trainer = GRPOTrainer(\n",
    "    model=model,\n",
    "    processing_class=tokenizer,\n",
    "    reward_funcs=[\n",
    "        xmlcount_reward_func,\n",
    "        soft_format_reward_func,\n",
    "        strict_format_reward_func,\n",
    "        int_reward_func,\n",
    "        accuracy_reward_func],\n",
    "    args=training_args,\n",
    "    train_dataset=dataset,\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
